#include <ros/ros.h>

#include <pcl_ros/point_cloud.h>
#include <pcl_ros/transforms.h>
#include <pcl/point_types.h>
#include <pcl/common/transforms.h>
#include <pcl/common/common.h>

#include <grid_map_ros/grid_map_ros.hpp>
#include <grid_map_core/GridMap.hpp>

#include <tf2_ros/transform_listener.h>
#include <tf2/LinearMath/Transform.h>
#include <tf2_eigen/tf2_eigen.h>
#include <tf2_geometry_msgs/tf2_geometry_msgs.h>
#include <tf/transform_broadcaster.h>

class PcdToElevation
{
public:
    PcdToElevation(ros::NodeHandle &nh) : nh_(nh)
    {
        ros::NodeHandle private_nh("~");

        private_nh.param<std::string>("point_cloud_topic", point_cloud_topic_, "/cloud_registered");
        private_nh.param<std::string>("elevation_topic", elevation_topic_, "/elevation_map");
        private_nh.param<std::string>("world_frame_id", world_frame_id_, "map");
        private_nh.param<std::string>("body_frame_id", body_frame_id_, "body");
        private_nh.param<double>("filter_size", filter_size_, 0.1);
        private_nh.param<double>("resolution", resolution_, 0.05);
        private_nh.param<double>("max_height", max_height_, 1.0);
        private_nh.param<int>("max_cell_num", max_cell_num_, 100000000); // Default: 100 million cells

        body_frame_2d_id_ = body_frame_id_ + "_2d";

        sub_ = nh_.subscribe(point_cloud_topic_, 1, &PcdToElevation::cloudCallback, this);
        pub_ = nh_.advertise<grid_map_msgs::GridMap>(elevation_topic_, 1);

        tf_timer_ = nh_.createTimer(ros::Duration(0.05), &PcdToElevation::publish2DTransform, this);

        tf_buffer_.reset(new tf2_ros::Buffer(ros::Duration(10.0)));
        tf_listener_.reset(new tf2_ros::TransformListener(*tf_buffer_));

        tf_broadcaster_.reset(new tf::TransformBroadcaster());
        if (!tf_broadcaster_)
        {
            ROS_FATAL("Failed to initialize TransformBroadcaster.");
            ros::shutdown();
        }

        ROS_INFO("PCD_To_Elevation node started. Max cells: %d", max_cell_num_);
    }

private:
    ros::NodeHandle nh_;
    ros::Subscriber sub_;
    ros::Publisher pub_;
    ros::Timer tf_timer_;

    std::string point_cloud_topic_;
    std::string world_frame_id_;
    std::string body_frame_id_;
    std::string elevation_topic_;
    std::string body_frame_2d_id_;
    double filter_size_;
    double resolution_;
    double max_height_;
    int max_cell_num_;

    std::shared_ptr<tf2_ros::Buffer> tf_buffer_;
    std::unique_ptr<tf2_ros::TransformListener> tf_listener_;
    std::unique_ptr<tf::TransformBroadcaster> tf_broadcaster_;

    void cloudCallback(const pcl::PointCloud<pcl::PointXYZ>::ConstPtr &cloud)
    {
        pcl::PointCloud<pcl::PointXYZ>::Ptr transformed_cloud(new pcl::PointCloud<pcl::PointXYZ>);
        pcl::PointCloud<pcl::PointXYZ>::ConstPtr input_cloud;
        if (cloud->header.frame_id != body_frame_2d_id_)
        {
            try
            {
                geometry_msgs::TransformStamped transform_stamped = tf_buffer_->lookupTransform(
                    body_frame_2d_id_, cloud->header.frame_id, ros::Time(0), ros::Duration(0.5));
                Eigen::Affine3f transform = tf2::transformToEigen(transform_stamped.transform).cast<float>();
                pcl::transformPointCloud(*cloud, *transformed_cloud, transform);
                input_cloud = transformed_cloud;
            }
            catch (tf2::TransformException &ex)
            {
                ROS_WARN("Failed to transform point cloud: %s", ex.what());
                return;
            }
        }
        else
        {
            input_cloud = cloud;
        }

        // Skip processing if cloud is empty
        if (input_cloud->empty())
        {
            // ROS_WARN("Input cloud is empty. Skipping.");
            return;
        }

        grid_map::GridMap elevation_map;
        elevation_map.setFrameId(world_frame_id_);

        Eigen::Vector4f min_pt_eigen, max_pt_eigen;
        pcl::getMinMax3D(*input_cloud, min_pt_eigen, max_pt_eigen);

        double y_min = std::min(static_cast<double>(min_pt_eigen[1]), 0.0);
        double y_max = std::max(static_cast<double>(max_pt_eigen[1]), 0.0);

        double length_x = max_pt_eigen[0] - min_pt_eigen[0] + 0.2;
        double length_y = y_max - y_min + 0.2;

        // Check grid size to prevent excessive memory allocation
        double cell_count = (length_x / resolution_) * (length_y / resolution_);
        if (cell_count > max_cell_num_)
        {
            ROS_WARN("Grid size too large (%.0f cells > %d). Skipping.", cell_count, max_cell_num_);
            return;
        }

        grid_map::Position position((min_pt_eigen[0] + max_pt_eigen[0]) / 2.0, 0.0);

        try
        {
            // Critical section: grid map allocation and processing
            elevation_map.setGeometry(grid_map::Length(length_x, length_y), resolution_, position);
            elevation_map.add("elevation", -std::numeric_limits<float>::infinity());

            for (const auto &point : input_cloud->points)
            {
                if (std::isnan(point.x) || std::isnan(point.y) || std::isnan(point.z))
                    continue;
                if (point.z > max_height_)
                    continue;

                grid_map::Position pos(point.x, point.y);
                grid_map::Index index;
                if (!elevation_map.getIndex(pos, index))
                    continue;

                index(0) = std::clamp(index(0), 0, elevation_map.getSize()(0) - 1);
                index(1) = std::clamp(index(1), 0, elevation_map.getSize()(1) - 1);

                if (elevation_map.at("elevation", index) < point.z)
                {
                    elevation_map.at("elevation", index) = point.z;
                }
            }

            // fillMissingElevationValues(elevation_map);

            grid_map_msgs::GridMap message;
            grid_map::GridMapRosConverter::toMessage(elevation_map, message);
            message.info.header.stamp = ros::Time::now();
            message.info.header.frame_id = body_frame_2d_id_;
            pub_.publish(message);
        }
        catch (const std::bad_alloc &e)
        {
            ROS_ERROR("Memory allocation failed: %s", e.what());
        }
        catch (const std::exception &e)
        {
            ROS_ERROR("Exception during grid map processing: %s", e.what());
        }
    }
    void fillMissingElevationValues(grid_map::GridMap &map)
    {
        const std::string layer = "elevation";
        const int rows = map.getSize()(0);
        const int cols = map.getSize()(1);
        const int search_radius = resolution_ * 500;

        for (int i = 0; i < rows; ++i)
        {
            for (int j = 0; j < cols; ++j)
            {
                grid_map::Index index(i, j);
                if (!std::isfinite(map.at(layer, index)))
                {
                    float weight_sum = 0.0f;
                    float z_sum = 0.0f;

                    for (int di = -search_radius; di <= search_radius; ++di)
                    {
                        for (int dj = -search_radius; dj <= search_radius; ++dj)
                        {
                            if (di == 0 && dj == 0)
                                continue;

                            grid_map::Index neighbor(i + di, j + dj);
                            if (map.isValid(neighbor))
                            {
                                float z = map.at(layer, neighbor);
                                if (std::isfinite(z))
                                {
                                    float dist = std::sqrt(di * di + dj * dj);
                                    float weight = 1.0f / (dist + 1e-6);
                                    weight_sum += weight;
                                    z_sum += z * weight;
                                }
                            }
                        }
                    }

                    if (weight_sum > 0.0f)
                    {
                        map.at(layer, index) = z_sum / weight_sum;
                    }
                    else
                    {
                        map.at(layer, index) = 0.0f;
                    }
                }
            }
        }
    }

    void publish2DTransform(const ros::TimerEvent &)
    {
        try
        {
            geometry_msgs::TransformStamped transform_stamped = tf_buffer_->lookupTransform(
                world_frame_id_, body_frame_id_, ros::Time(0), ros::Duration(0.1));

            tf::Vector3 origin;
            tf::vector3MsgToTF(transform_stamped.transform.translation, origin);

            tf::Quaternion quat;
            tf::quaternionMsgToTF(transform_stamped.transform.rotation, quat);
            double roll, pitch, yaw;
            tf::Matrix3x3(quat).getEulerYPR(yaw, pitch, roll);

            tf::Quaternion new_quat;
            new_quat.setRPY(0, 0, yaw);

            if (std::isnan(origin.x()) || std::isnan(origin.y()) || std::isnan(origin.z()))
            {
                ROS_WARN("Invalid origin for TF transform.");
                return;
            }
            if (std::isnan(new_quat.x()) || std::isnan(new_quat.y()) ||
                std::isnan(new_quat.z()) || std::isnan(new_quat.w()))
            {
                ROS_WARN("Invalid quaternion for TF transform.");
                return;
            }

            tf::Transform transform(new_quat, origin);
            tf::StampedTransform stamped_transform(
                transform, ros::Time::now(), world_frame_id_, body_frame_2d_id_.c_str());
            tf_broadcaster_->sendTransform(stamped_transform);
        }
        catch (tf2::TransformException &ex)
        {
            ROS_WARN("Failed to get transform for 2D frame: %s", ex.what());
        }
    }
};

int main(int argc, char **argv)
{
    ros::init(argc, argv, "pcd_to_elevation_node");
    ros::NodeHandle nh;
    PcdToElevation node(nh);
    ros::spin();
    return 0;
}
